# -*- coding: utf-8 -*-
"""
A股资金流策略
基于北向资金、主力大单和融资融券数据捕捉资金流向，包含A股特有机制处理
"""

from zipline.api import order_value, symbol, get_datetime
from strategies.utils.a_share_utils import get_csi300_constituents, is_price_limit, apply_t1_restriction
from zipline.finance import commission, slippage
import pandas as pd
import numpy as np

def initialize(context):
    """
    初始化策略
    """
    # 设置策略参数
    context.north_weight = 0.4      # 北向资金权重
    context.big_order_weight = 0.4  # 大单资金权重
    context.margin_weight = 0.2     # 两融资金权重
    context.flow_threshold = 1e8    # 流入阈值 (1亿人民币)
    context.max_position = 0.05     # 单票最大仓位
    
    # 设置交易费用 (A股标准)
    context.set_commission(commission.PerShare(cost=0.0003, min_trade_cost=5))
    context.set_slippage(slippage.FixedSlippage(spread=0.001))

def handle_data(context, data):
    """
    每日交易逻辑
    """
    # 获取沪深300成分股
    stocks = get_csi300_constituents(context, data)
    
    for stock in stocks:
        # 跳过涨跌停股票
        if is_price_limit(stock, data):
            continue
            
        # 获取资金流数据
        north_flow = get_northbound_flow(stock)
        big_order_flow = get_big_order_flow(stock)
        margin_flow = get_margin_flow(stock)
        
        # 计算综合资金流得分
        flow_score = (
            context.north_weight * north_flow +
            context.big_order_weight * big_order_flow +
            context.margin_weight * margin_flow
        )
        
        # 买入信号
        if flow_score > context.flow_threshold:
            # 应用T+1限制
            closeable_amount = apply_t1_restriction(context, stock)
            if closeable_amount == 0:  # 无持仓时可买入
                order_value(stock, context.portfolio.cash * context.max_position)

def get_northbound_flow(stock):
    """
    获取北向资金净流入 (模拟)
    """
    # 在实际使用中应替换为真实数据接口
    np.random.seed(hash(stock) % 1000)
    return np.random.uniform(-5e7, 5e7)  # 随机生成-5000万到5000万的资金流

def get_big_order_flow(stock):
    """
    获取主力大单净流入 (模拟)
    """
    # 在实际使用中应替换为真实数据接口
    np.random.seed(hash(stock) % 2000)
    return np.random.uniform(-1e8, 1e8)  # 随机生成-1亿到1亿的资金流

def get_margin_flow(stock):
    """
    获取融资融券净流入 (模拟)
    """
    # 在实际使用中应替换为真实数据接口
    np.random.seed(hash(stock) % 3000)
    return np.random.uniform(-5e7, 5e7)  # 随机生成-5000万到5000万的资金流

def analyze(context, perf):
    """
    回测结果分析
    """
    # 输出策略表现
    print("策略年化收益率: %.2f%%" % (perf.returns.mean() * 252 * 100))
    print("最大回撤: %.2f%%" % (perf.max_drawdown() * 100))
    print("夏普比率: %.2f" % perf.sharpe_ratio)
    
    # 绘制净值曲线
    perf.portfolio_value.plot()